﻿using Fst.SemiRing;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Util;

namespace Fst.Operations
{
    /// <summary>
    /// N-shortest paths operation.
    /// 
    /// See: M. Mohri, M. Riley,
    /// "An Efficient Algorithm for the n-best-strings problem", Proceedings of the
    /// International Conference on Spoken Language Processing 2002 (ICSLP '02).
    /// 
    /// See: M. Mohri,
    /// "Semiring Framework and Algorithms for Shortest-Distance Problems", Journal
    /// of Automata, Languages and Combinatorics, 7(3), pp. 321-350, 2002.
    /// 
    /// @author John Salatas <jsalatas@users.sourceforge.net>
    /// </summary>
    public class NShortestPaths
    {
        private NShortestPaths() {
        }

        /**
        /// Calculates the shortest distances from each state to the final
        /// 
        /// @param fst
        ///            the fst to calculate the shortest distances
        /// @return the array containing the shortest distances
         */
        public static float[] shortestDistance(Fst fst) 
        {

            Fst reversed = Reverse.get(fst);

            float[] d = new float[reversed.getNumStates()];
            float[] r = new float[reversed.getNumStates()];

            Semiring semiring = reversed.getSemiring();

            Array.ForEach(d, x => semiring.zero());
            Array.ForEach(r, x => semiring.zero());

            List<State> queue = new List<State>();

            queue.Add(reversed.getStart());

            d[reversed.getStart().getId()] = semiring.one();
            r[reversed.getStart().getId()] = semiring.one();

            while (queue.Count!=0) 
            {
                State q = queue.First();
                queue.Remove(q);

                float rnew = r[q.getId()];
                r[q.getId()] = semiring.zero();

                for (int i = 0; i < q.getNumArcs(); i++) {
                    Arc a = q.getArc(i);
                    State nextState = a.getNextState();
                    float dnext = d[a.getNextState().getId()];
                    float dnextnew = semiring.plus(dnext,
                            semiring.times(rnew, a.getWeight()));
                    if (dnext != dnextnew) {
                        d[a.getNextState().getId()] = dnextnew;
                        r[a.getNextState().getId()] = semiring.plus(r[a
                                .getNextState().getId()], semiring.times(rnew,
                                a.getWeight()));
                        if (!queue.Contains(nextState)) 
                        {
                            queue.Add(nextState);
                        }
                    }
                }
            }
            return d;
        }

        /**
        /// Calculates the n-best shortest path from the initial to the final state.
        /// 
        /// @param fst
        ///            the fst to calculate the nbest shortest paths
        /// @param n
        ///            number of best paths to return
        /// @param determinize
        ///            if true the input fst will bwe determinized prior the
        ///            operation
        /// @return an fst containing the n-best shortest paths
         */
        public static Fst get(Fst fst, int n, Boolean determinize) 
        {
            if (fst == null) 
            {
                return null;
            }

            if (fst.getSemiring() == null) {
                return null;
            }
            Fst fstdet = fst;
            if (determinize) {
                fstdet = Determinize.get(fst);
            }
            Semiring semiring = fstdet.getSemiring();
            Fst res = new Fst(semiring);
            res.setIsyms(fstdet.getIsyms());
            res.setOsyms(fstdet.getOsyms());

            float[] d = shortestDistance(fstdet);

            ExtendFinal.apply(fstdet);

            int[] r = new int[fstdet.getNumStates()];

            SortedSet<Pair<State, float>> queue = new SortedSet<Pair<State, float>>();
            //x =>
            //        {
            //            public int compare(Pair<State, float> o1,
            //                    Pair<State, float> o2) {
            //                float previous = o1.getRight();
            //                float d1 = d[o1.getLeft().getId()];

            //                float next = o2.getRight();
            //                float d2 = d[o2.getLeft().getId()];

            //                float a1 = semiring.times(next, d2);
            //                float a2 = semiring.times(previous, d1);

            //                if (semiring.naturalLess(a1, a2))
            //                    return 1;

            //                if (a1 == a2)
            //                    return 0;

            //                return -1;
            //            }
            //        });

            SortedList<Pair<State, float>, Pair<State, float>> previous = new SortedList<Pair<State, float>, Pair<State, float>>(
                    fst.getNumStates());
            SortedList<Pair<State, float>, State> stateMap = new SortedList<Pair<State, float>, State>(
                    fst.getNumStates());

            State start = fstdet.getStart();
            Pair<State, float> item = new Pair<State, float>(start, semiring.one());
            queue.Add(item);
            previous.Add(item, null);

            while (queue.Count!=0) 
            {
                Pair<State, float> pair = queue.First();
                queue.Remove(pair);
                State p = pair.getLeft();
                float c = pair.getRight();

                State s = new State(p.getFinalWeight());
                res.addState(s);
                stateMap.Add(pair, s);
                if (previous[pair] == null) {
                    // this is the start state
                    res.setStart(s);
                } else {
                    // add the incoming arc from previous to current
                    State previouState = stateMap[previous[pair]];
                    State previousOldState = previous[pair].getLeft();
                    for (int j = 0; j < previousOldState.getNumArcs(); j++) {
                        Arc a = previousOldState.getArc(j);
                        if (a.getNextState().Equals(p)) 
                        {
                            previouState.addArc(new Arc(a.getIlabel(), a
                                    .getOlabel(), a.getWeight(), s));
                        }
                    }
                }

                int stateIndex = p.getId();
                r[stateIndex]++;

                if ((r[stateIndex] == n) && (p.getFinalWeight() != semiring.zero())) {
                    break;
                }

                if (r[stateIndex] <= n) {
                    for (int j = 0; j < p.getNumArcs(); j++) {
                        Arc a = p.getArc(j);
                        float cnew = semiring.times(c, a.getWeight());
                        Pair<State, float> next = new Pair<State, float>(
                                a.getNextState(), cnew);
                        previous.Add(next, pair);
                        queue.Add(next);
                    }
                }
            }

            return res;
        }
    }
}
